"""
定时任务调度服务
基于APScheduler实现定时任务管理
"""

import uuid
from datetime import datetime
from typing import Dict, List, Optional, Any
from apscheduler.schedulers.background import BackgroundScheduler
from apscheduler.jobstores.sqlalchemy import SQLAlchemyJobStore
from apscheduler.triggers.cron import CronTrigger
from apscheduler.triggers.interval import IntervalTrigger
from apscheduler.triggers.date import DateTrigger
from apscheduler.events import EVENT_JOB_EXECUTED, EVENT_JOB_ERROR
import sqlalchemy as sa
from sqlalchemy.orm import sessionmaker
import logging

from src.config.settings import settings
from src.services.database_session_service import DatabaseSessionService
from src.research_core.enhanced_multi_agent_workflow import create_enhanced_multi_agent_workflow as create_multi_agent_workflow

logger = logging.getLogger(__name__)

class SchedulerService:
    """定时任务调度服务"""
    
    def __init__(self):
        self.scheduler = None
        self.workflow = None
        self.engine = None
        self.SessionLocal = None
        self.session_service = None
        
        # 立即初始化必要组件
        self.initialize_database()
        
    def initialize_database(self):
        """初始化数据库连接"""
        try:
            from src.config.settings import settings
            # 初始化数据库连接
            self.engine = sa.create_engine(settings.DATABASE_URL)
            self.SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=self.engine)
            
            # 初始化session_service
            self.session_service = DatabaseSessionService(self.db_session_factory())
        except Exception as e:
            logger.error(f"初始化数据库失败: {e}", exc_info=True)
            # 即使初始化失败，也要继续运行，以便其他功能可以正常工作
            
    def db_session_factory(self):
        """创建数据库会话的工厂方法"""
        if not self.SessionLocal:
            raise RuntimeError("数据库会话工厂尚未初始化")
        return self.SessionLocal()
        
    def initialize(self):
        """初始化调度器"""
        try:
            # 确保数据库已初始化
            if not self.engine or not self.SessionLocal:
                self.initialize_database()
            
            # 配置作业存储
            jobstores = {
                'default': SQLAlchemyJobStore(
                    engine=self.engine,
                    tablename='apscheduler_jobs'
                )
            }
            
            # 创建调度器
            self.scheduler = BackgroundScheduler(
                jobstores=jobstores,
                coalesce=settings.SCHEDULER_COALESCE,
                max_instances=settings.SCHEDULER_MAX_INSTANCES
            )
            
            # 添加事件监听器
            self.scheduler.add_listener(self._job_executed, EVENT_JOB_EXECUTED)
            self.scheduler.add_listener(self._job_error, EVENT_JOB_ERROR)
            
            # 初始化工作流
            self.workflow = create_multi_agent_workflow().compile()
            
            logger.info("定时任务调度器初始化成功")
            
        except Exception as e:
            logger.error(f"定时任务调度器初始化失败: {e}")
            raise
    
    def _job_executed(self, event):
        """任务执行成功事件处理"""
        logger.info(f"任务执行成功: {event.job_id}")
    
    def _job_error(self, event):
        """任务执行失败事件处理"""
        logger.error(f"任务执行失败: {event.job_id}, 错误: {event.exception}")
    
    def _execute_research_task(self, question: str, session_id: Optional[str] = None):
        """执行研究任务"""
        try:
            if not session_id:
                session_id = self.session_service.create_session(f"定时任务_{datetime.now().strftime('%Y%m%d_%H%M%S')}")
            
            # 定义初始状态
            initial_state = {
                "question": question,
                "search_query": "",
                "search_results": "",
                "research_complete": False,
                "final_answer": "",
                "search_strategy": "",
                "analysis_results": {},
                "agent_assignments": {}
            }
            
            # 运行工作流
            final_state = self.workflow.invoke(initial_state)
            
            logger.info(f"定时任务执行完成: {question}")
            return {
                "success": True,
                "session_id": session_id,
                "answer": final_state['final_answer'],
                "research_complete": final_state.get('research_complete', False)
            }
            
        except Exception as e:
            logger.error(f"定时任务执行失败: {e}")
            return {
                "success": False,
                "error": str(e)
            }
    
    def add_cron_task(self, task_id: str, question: str, cron_expression: str, 
                     session_id: Optional[str] = None, **kwargs) -> bool:
        """添加cron定时任务"""
        try:
            self.scheduler.add_job(
                self._execute_research_task,
                trigger=CronTrigger.from_crontab(cron_expression),
                id=task_id,
                args=[question],
                kwargs={'session_id': session_id},
                **kwargs
            )
            logger.info(f"添加cron定时任务成功: {task_id}")
            return True
        except Exception as e:
            logger.error(f"添加cron定时任务失败: {e}")
            return False
    
    def add_interval_task(self, task_id: str, question: str, interval_seconds: int,
                         session_id: Optional[str] = None, **kwargs) -> bool:
        """添加间隔定时任务"""
        try:
            self.scheduler.add_job(
                self._execute_research_task,
                trigger=IntervalTrigger(seconds=interval_seconds),
                id=task_id,
                args=[question],
                kwargs={'session_id': session_id},
                **kwargs
            )
            logger.info(f"添加间隔定时任务成功: {task_id}")
            return True
        except Exception as e:
            logger.error(f"添加间隔定时任务失败: {e}")
            return False
    
    def add_date_task(self, task_id: str, question: str, run_date: datetime,
                     session_id: Optional[str] = None, **kwargs) -> bool:
        """添加一次性定时任务"""
        try:
            self.scheduler.add_job(
                self._execute_research_task,
                trigger=DateTrigger(run_date=run_date),
                id=task_id,
                args=[question],
                kwargs={'session_id': session_id},
                **kwargs
            )
            logger.info(f"添加一次性定时任务成功: {task_id}")
            return True
        except Exception as e:
            logger.error(f"添加一次性定时任务失败: {e}")
            return False
    
    def remove_task(self, task_id: str) -> bool:
        """移除定时任务"""
        try:
            self.scheduler.remove_job(task_id)
            logger.info(f"移除定时任务成功: {task_id}")
            return True
        except Exception as e:
            logger.error(f"移除定时任务失败: {e}")
            return False
    
    def pause_task(self, task_id: str) -> bool:
        """暂停定时任务"""
        try:
            self.scheduler.pause_job(task_id)
            logger.info(f"暂停定时任务成功: {task_id}")
            return True
        except Exception as e:
            logger.error(f"暂停定时任务失败: {e}")
            return False
    
    def resume_task(self, task_id: str) -> bool:
        """恢复定时任务"""
        try:
            self.scheduler.resume_job(task_id)
            logger.info(f"恢复定时任务成功: {task_id}")
            return True
        except Exception as e:
            logger.error(f"恢复定时任务失败: {e}")
            return False
    
    def get_task(self, task_id: str) -> Optional[Dict]:
        """获取定时任务信息"""
        try:
            job = self.scheduler.get_job(task_id)
            if job:
                return {
                    "id": job.id,
                    "name": job.name,
                    "next_run_time": job.next_run_time,
                    "trigger": str(job.trigger),
                    "args": job.args,
                    "kwargs": job.kwargs
                }
            return None
        except Exception as e:
            logger.error(f"获取定时任务信息失败: {e}")
            return None
    
    def list_tasks(self) -> List[Dict]:
        """获取所有定时任务"""
        try:
            jobs = self.scheduler.get_jobs()
            return [
                {
                    "id": job.id,
                    "name": job.name,
                    "next_run_time": job.next_run_time,
                    "trigger": str(job.trigger),
                    "args": job.args,
                    "kwargs": job.kwargs
                }
                for job in jobs
            ]
        except Exception as e:
            logger.error(f"获取定时任务列表失败: {e}")
            return []
    
    def start(self):
        """启动调度器"""
        if self.scheduler and not self.scheduler.running:
            self.scheduler.start()
            logger.info("定时任务调度器已启动")
    
    def shutdown(self):
        """关闭调度器"""
        if self.scheduler and self.scheduler.running:
            self.scheduler.shutdown()
            logger.info("定时任务调度器已关闭")

    def get_execution_history(self, db, task_id: str, limit: int = 10) -> list:
        """获取任务执行历史"""
        try:
            from src.models.scheduled_task import TaskExecutionHistory
            history = db.query(TaskExecutionHistory)\
                .filter(TaskExecutionHistory.task_id == task_id)\
                .order_by(TaskExecutionHistory.started_at.desc())\
                .limit(limit)\
                .all()
            return history
        except Exception as e:
            logger.error(f"获取执行历史记录失败: {e}")
            return []

    def create_task(self, db, task_data: Dict[str, Any]):
        """创建定时任务"""
        try:
            from src.models.scheduled_task import ScheduledTask
            task = ScheduledTask(**task_data)
            db.add(task)
            db.commit()
            db.refresh(task)
            return task
        except Exception as e:
            logger.error(f"创建定时任务失败: {e}")
            db.rollback()
            return None

    def get_task(self, db, task_id: str):
        """获取定时任务"""
        try:
            from src.models.scheduled_task import ScheduledTask
            task = db.query(ScheduledTask).filter(ScheduledTask.id == task_id).first()
            return task
        except Exception as e:
            logger.error(f"获取定时任务失败: {e}")
            return None

    def update_task(self, db, task_id: str, update_data: Dict[str, Any]) -> bool:
        """更新定时任务"""
        try:
            from src.models.scheduled_task import ScheduledTask
            task = db.query(ScheduledTask).filter(ScheduledTask.id == task_id).first()
            if task:
                for key, value in update_data.items():
                    setattr(task, key, value)
                db.commit()
                return True
            return False
        except Exception as e:
            logger.error(f"更新定时任务失败: {e}")
            db.rollback()
            return False

    def delete_task(self, db, task_id: str) -> bool:
        """删除定时任务"""
        try:
            from src.models.scheduled_task import ScheduledTask
            task = db.query(ScheduledTask).filter(ScheduledTask.id == task_id).first()
            if task:
                db.delete(task)
                db.commit()
                return True
            return False
        except Exception as e:
            logger.error(f"删除定时任务失败: {e}")
            db.rollback()
            return False

    def list_tasks(self, db, enabled_only: bool = False) -> list:
        """获取定时任务列表"""
        try:
            from src.models.scheduled_task import ScheduledTask
            query = db.query(ScheduledTask)
            if enabled_only:
                query = query.filter(ScheduledTask.enabled == True)
            tasks = query.all()
            return tasks
        except Exception as e:
            logger.error(f"获取定时任务列表失败: {e}")
            return []

# 全局调度器实例
scheduler_service = SchedulerService()
